import os
import sys
import argparse
import json
import re
import uuid
from pdf2image import convert_from_path
import cv2
import numpy as np
from paddleocr import PaddleOCR

class PDFProcessor:
    def __init__(self, lang='ch'):
        # 初始化OCR引擎
        self.ocr = PaddleOCR(use_angle_cls=True, lang=lang)
        print("初始化OCR引擎完成")
    
    def process_pdf_directory(self, pdf_dir, output_dir):
        """处理目录中的所有PDF文件"""
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)
            
        # 获取所有PDF文件
        pdf_files = [f for f in os.listdir(pdf_dir) if f.lower().endswith('.pdf')]
        
        if not pdf_files:
            print(f"在 {pdf_dir} 中没有找到PDF文件")
            return
        
        print(f"找到 {len(pdf_files)} 个PDF文件，开始处理...")
        
        for pdf_file in pdf_files:
            pdf_path = os.path.join(pdf_dir, pdf_file)
            print(f"处理: {pdf_file}")
            self.process_pdf(pdf_path, output_dir)
    
    def sanitize_filename(self, filename):
        """清理文件名，移除特殊字符"""
        # 移除不安全的字符，只保留字母、数字、下划线和连字符
        safe_name = re.sub(r'[^\w\-]', '_', filename)
        # 如果文件名太长，截断它
        if len(safe_name) > 50:
            safe_name = safe_name[:50]
        return safe_name
    
    def process_pdf(self, pdf_path, output_dir):
        """处理单个PDF文件"""
        # 创建PDF文件的输出目录，使用安全的文件名
        pdf_name = os.path.basename(pdf_path).split('.')[0]
        safe_pdf_name = self.sanitize_filename(pdf_name)
        
        # 为确保唯一性，添加一个短UUID
        unique_id = str(uuid.uuid4())[:8]
        pdf_output_dir = os.path.join(output_dir, f"{safe_pdf_name}_{unique_id}")
        os.makedirs(pdf_output_dir, exist_ok=True)
        
        # 将PDF转换为图像
        try:
            images = convert_from_path(pdf_path)
        except Exception as e:
            print(f"转换PDF时出错: {e}")
            return
        
        print(f"PDF共 {len(images)} 页，开始处理每一页...")
        
        # 处理每一页
        results = []
        for i, image in enumerate(images):
            # 保存图像
            img_path = os.path.join(pdf_output_dir, f"page_{i+1}.png")
            image.save(img_path, "PNG")
            
            # 验证图像是否成功保存
            if not os.path.exists(img_path):
                print(f"警告: 图像未能保存到 {img_path}")
                continue
                
            try:
                # 处理图像
                page_result = self.process_image(img_path, i+1)
                results.append(page_result)
                print(f"  页面 {i+1}/{len(images)} 处理完成")
            except Exception as e:
                print(f"  处理页面 {i+1} 时出错: {e}")
                # 继续处理下一页
                continue
        
        # 保存整个PDF的结果
        output_json = os.path.join(output_dir, f"{safe_pdf_name}_{unique_id}_result.json")
        with open(output_json, 'w', encoding='utf-8') as f:
            json.dump(results, f, ensure_ascii=False, indent=2)
        
        print(f"PDF处理完成，结果保存至: {output_json}")
        return output_json
    
    def process_image(self, img_path, page_num):
        """处理单个图像，提取文本和布局"""
        # 读取图像
        image = cv2.imread(img_path)
        if image is None:
            print(f"尝试使用其他方法读取图像: {img_path}")
            # 尝试使用PIL读取
            from PIL import Image
            try:
                pil_image = Image.open(img_path)
                # 转换为OpenCV格式
                image = cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR)
            except Exception as e:
                raise ValueError(f"无法读取图像: {img_path}, 错误: {e}")
        
        # OCR识别
        ocr_result = self.ocr.ocr(img_path, cls=True)
        
        # 简化版布局分析
        height, width = image.shape[:2]
        
        # 构建结果
        result = {
            "page": page_num,
            "image": os.path.basename(img_path),
            "blocks": []
        }
        
        # 如果OCR结果为空，返回空结果
        if not ocr_result[0]:
            return result
        
        # 提取所有文本
        all_texts = []
        for line in ocr_result[0]:
            bbox = line[0]
            text = line[1][0]
            confidence = line[1][1]
            all_texts.append({
                "text": text,
                "confidence": float(confidence),
                "bbox": [[float(p[0]), float(p[1])] for p in bbox]
            })
        
        # 检测表格
        table_candidates = self.detect_tables_heuristic(all_texts, width, height)
        
        # 处理表格
        for i, table_coords in enumerate(table_candidates):
            x1, y1, x2, y2 = table_coords
            
            # 提取表格内的文本
            table_texts = []
            for text in all_texts:
                text_x1 = text["bbox"][0][0]
                text_y1 = text["bbox"][0][1]
                if (x1 <= text_x1 <= x2 and y1 <= text_y1 <= y2):
                    table_texts.append(text)
            
            # 提取表格结构
            table_content = self.extract_table_structure(table_texts)
            
            result["blocks"].append({
                "type": "table",
                "coordinates": [int(x1), int(y1), int(x2), int(y2)],
                "content": table_content
            })
        
        # 剩余文本作为普通文本块
        remaining_texts = []
        for text in all_texts:
            is_in_table = False
            for table in table_candidates:
                text_x1 = text["bbox"][0][0]
                text_y1 = text["bbox"][0][1]
                if (table[0] <= text_x1 <= table[2] and table[1] <= text_y1 <= table[3]):
                    is_in_table = True
                    break
            
            if not is_in_table:
                remaining_texts.append(text)
        
        if remaining_texts:
            # 按Y坐标分组为段落
            paragraphs = self.group_texts_into_paragraphs(remaining_texts)
            
            for i, para_texts in enumerate(paragraphs):
                # 计算段落边界
                min_x = min(text["bbox"][0][0] for text in para_texts)
                min_y = min(text["bbox"][0][1] for text in para_texts)
                max_x = max(text["bbox"][2][0] for text in para_texts)
                max_y = max(text["bbox"][2][1] for text in para_texts)
                
                # 判断是否为标题 (简单启发式：短且位于顶部)
                is_title = (len(para_texts) <= 2 and min_y < height * 0.2)
                
                result["blocks"].append({
                    "type": "title" if is_title else "text",
                    "coordinates": [int(min_x), int(min_y), int(max_x), int(max_y)],
                    "content": [text["text"] for text in para_texts]
                })
        
        return result
    
    def detect_tables_heuristic(self, texts, width, height):
        """使用启发式方法检测表格"""
        if not texts:
            return []
            
        # 按Y坐标排序
        texts_by_y = sorted(texts, key=lambda x: x["bbox"][0][1])
        
        # 分组为行
        rows = []
        current_row = [texts_by_y[0]]
        row_y = texts_by_y[0]["bbox"][0][1]
        
        for text in texts_by_y[1:]:
            text_y = text["bbox"][0][1]
            # 如果Y坐标相差不大，认为是同一行
            if abs(text_y - row_y) < 20:  # 阈值可调整
                current_row.append(text)
            else:
                # 按X坐标排序当前行
                current_row.sort(key=lambda x: x["bbox"][0][0])
                rows.append(current_row)
                # 开始新行
                current_row = [text]
                row_y = text_y
        
        # 处理最后一行
        if current_row:
            current_row.sort(key=lambda x: x["bbox"][0][0])
            rows.append(current_row)
        
        # 检测表格特征
        table_candidates = []
        
        # 至少需要3行才考虑是表格
        if len(rows) >= 3:
            # 检查行是否有相似的列数和间距
            row_lengths = [len(row) for row in rows]
            
            # 如果多行有相同的列数，可能是表格
            for i in range(len(rows) - 2):
                if row_lengths[i] >= 2 and row_lengths[i] == row_lengths[i+1] == row_lengths[i+2]:
                    # 找到可能的表格，计算边界
                    table_rows = rows[i:i+3]
                    min_x = min(text["bbox"][0][0] for row in table_rows for text in row)
                    min_y = min(text["bbox"][0][1] for row in table_rows for text in row)
                    max_x = max(text["bbox"][2][0] for row in table_rows for text in row)
                    max_y = max(text["bbox"][2][1] for row in table_rows for text in row)
                    
                    # 扩展表格边界，包含后续相似行
                    j = i + 3
                    while j < len(rows) and row_lengths[j] == row_lengths[i]:
                        for text in rows[j]:
                            min_x = min(min_x, text["bbox"][0][0])
                            min_y = min(min_y, text["bbox"][0][1])
                            max_x = max(max_x, text["bbox"][2][0])
                            max_y = max(max_y, text["bbox"][2][1])
                        j += 1
                    
                    # 添加边距
                    padding = 10
                    table_candidates.append([
                        max(0, min_x - padding),
                        max(0, min_y - padding),
                        min(width, max_x + padding),
                        min(height, max_y + padding)
                    ])
                    
                    # 跳过已处理的行
                    i = j - 1
        
        return table_candidates
    
    def extract_table_structure(self, texts):
        """提取表格结构"""
        if not texts:
            return []
            
        # 按Y坐标排序
        texts.sort(key=lambda x: x["bbox"][0][1])
        
        # 简单的行分组
        rows = []
        current_row = [texts[0]]
        row_y = texts[0]["bbox"][0][1]
        
        for text in texts[1:]:
            text_y = text["bbox"][0][1]
            # 如果Y坐标相差不大，认为是同一行
            if abs(text_y - row_y) < 20:  # 阈值可调整
                current_row.append(text)
            else:
                # 按X坐标排序当前行
                current_row.sort(key=lambda x: x["bbox"][0][0])
                rows.append([t["text"] for t in current_row])
                # 开始新行
                current_row = [text]
                row_y = text_y
        
        # 处理最后一行
        if current_row:
            current_row.sort(key=lambda x: x["bbox"][0][0])
            rows.append([t["text"] for t in current_row])
            
        return rows
    
    def group_texts_into_paragraphs(self, texts):
        """将文本分组为段落"""
        if not texts:
            return []
            
        # 按Y坐标排序
        texts_by_y = sorted(texts, key=lambda x: x["bbox"][0][1])
        
        # 分组为段落
        paragraphs = []
        current_para = [texts_by_y[0]]
        para_y = texts_by_y[0]["bbox"][2][1]  # 使用底部Y坐标
        
        for text in texts_by_y[1:]:
            text_y = text["bbox"][0][1]  # 使用顶部Y坐标
            # 如果新文本的顶部与当前段落底部距离较近，认为是同一段落
            if abs(text_y - para_y) < 30:  # 段落间距阈值
                current_para.append(text)
                # 更新段落底部坐标
                para_y = max(para_y, text["bbox"][2][1])
            else:
                # 按X坐标排序当前段落中的文本
                current_para.sort(key=lambda x: x["bbox"][0][0])
                paragraphs.append(current_para)
                # 开始新段落
                current_para = [text]
                para_y = text["bbox"][2][1]
        
        # 处理最后一个段落
        if current_para:
            current_para.sort(key=lambda x: x["bbox"][0][0])
            paragraphs.append(current_para)
            
        return paragraphs

def main():
    parser = argparse.ArgumentParser(description='PDF批处理OCR工具')
    parser.add_argument('--pdf_dir', default='pdffiles', help='PDF文件目录')
    parser.add_argument('--output_dir', default='output_folder', help='输出目录')
    parser.add_argument('--lang', default='ch', help='OCR语言')
    
    args = parser.parse_args()
    
    # 确保输入目录存在
    if not os.path.exists(args.pdf_dir):
        print(f"PDF目录 '{args.pdf_dir}' 不存在，请先创建该目录并放入PDF文件")
        return
        
    processor = PDFProcessor(lang=args.lang)
    processor.process_pdf_directory(args.pdf_dir, args.output_dir)

if __name__ == "__main__":
    main() 